#!/bin/bash
# 清理所有非best的checkpoint，只保留最佳模型
# 使用方法: bash cleanup_keep_best_only.sh [checkpoint_dir]

CHECKPOINT_DIR=${1:-"./model/uncond_checkpoints_t3"}

echo "=========================================="
echo "清理非best checkpoint"
echo "=========================================="
echo "目录: ${CHECKPOINT_DIR}"
echo ""

# 检查目录是否存在
if [ ! -d "${CHECKPOINT_DIR}" ]; then
    echo "❌ 错误: 目录不存在: ${CHECKPOINT_DIR}"
    exit 1
fi

# 统计当前文件
echo "当前文件统计:"
echo "  总checkpoint数: $(ls ${CHECKPOINT_DIR}/*.pt 2>/dev/null | wc -l)"
echo "  Best checkpoint: $(ls ${CHECKPOINT_DIR}/best_*.pt 2>/dev/null | wc -l)"
echo "  普通checkpoint: $(ls ${CHECKPOINT_DIR}/model*.pt ${CHECKPOINT_DIR}/ema_*.pt ${CHECKPOINT_DIR}/opt*.pt 2>/dev/null | grep -v best | wc -l)"
echo ""

# 计算空间
BEFORE_SIZE=$(du -sh ${CHECKPOINT_DIR} | cut -f1)
echo "  当前占用空间: ${BEFORE_SIZE}"
echo ""

# 询问确认
echo "⚠️  警告: 将删除所有非best的checkpoint!"
echo "将保留的文件:"
ls ${CHECKPOINT_DIR}/best_*.pt 2>/dev/null | while read f; do
    echo "  ✓ $(basename $f)"
done

if [ -z "$2" ] || [ "$2" != "--confirm" ]; then
    echo ""
    read -p "确认继续? (输入 yes 继续): " confirm
    if [ "$confirm" != "yes" ]; then
        echo "已取消"
        exit 0
    fi
fi

echo ""
echo "开始清理..."
echo ""

# 删除普通checkpoint
deleted=0

# 删除 model*.pt (非best)
for f in ${CHECKPOINT_DIR}/model[0-9]*.pt; do
    if [ -f "$f" ]; then
        rm "$f"
        echo "  ✗ 删除: $(basename $f)"
        ((deleted++))
    fi
done

# 删除 ema_*.pt (非best)
for f in ${CHECKPOINT_DIR}/ema_*[0-9].pt; do
    if [ -f "$f" ] && [[ ! "$f" =~ best_ ]]; then
        rm "$f"
        echo "  ✗ 删除: $(basename $f)"
        ((deleted++))
    fi
done

# 删除 opt*.pt (非best)
for f in ${CHECKPOINT_DIR}/opt[0-9]*.pt; do
    if [ -f "$f" ]; then
        rm "$f"
        echo "  ✗ 删除: $(basename $f)"
        ((deleted++))
    fi
done

echo ""
echo "=========================================="
echo "清理完成"
echo "=========================================="
echo "删除文件数: ${deleted}"
echo ""

# 最终统计
echo "清理后文件统计:"
echo "  总checkpoint数: $(ls ${CHECKPOINT_DIR}/*.pt 2>/dev/null | wc -l)"
echo "  Best checkpoint: $(ls ${CHECKPOINT_DIR}/best_*.pt 2>/dev/null | wc -l)"
echo ""

AFTER_SIZE=$(du -sh ${CHECKPOINT_DIR} | cut -f1)
echo "  清理后空间: ${AFTER_SIZE} (之前: ${BEFORE_SIZE})"
echo ""

# 显示保留的best文件
echo "保留的best checkpoint:"
ls -lh ${CHECKPOINT_DIR}/best_*.pt 2>/dev/null | while read line; do
    echo "  $line"
done

echo ""
echo "✅ 完成！只保留了最佳模型"

